import matplotlib.pyplot as plt
import numpy as np

# Create figure with two subplots
fig, (ax1, ax2) = plt.subplots(1, 2, figsize=(10, 5))

# ========== Subplot 1: Local Computation Time (Fig 7a) ==========
methods = ['FedAvg', 'FedProx', 'FedEntGate']
low_power = [320, 285, 180]  # Low-power devices (e.g., Raspberry Pi)
high_power = [220, 190, 140]  # High-power devices (e.g., GPU server)

x = np.arange(len(methods))
width = 0.35

# Plot bars
patterns = ['/', 'x'] 
rects1 = ax1.bar(x - width/2, low_power, width, label='Low-power Devices', color='#1f77b4', hatch=patterns[0])
rects2 = ax1.bar(x + width/2, high_power, width, label='High-power Devices', color='#ff7f0e', hatch=patterns[1])

# Set labels and title
ax1.set_ylabel('Computation Time (ms)', fontsize=12)
ax1.set_title('(a) Local Computation Time Comparison (CIFAR-10)', fontsize=14,y=-0.3)
ax1.set_xticks(x)
ax1.set_xticklabels(methods)
ax1.legend(loc='upper right')
ax1.grid(axis='y', linestyle='--', alpha=0.7)

# Annotate optimization percentage
for rect in rects1:
    height = rect.get_height()
    if rect.get_x() > 1.5:  # FedEntGate position
        ax1.annotate(f'-28%', 
                    xy=(rect.get_x() + rect.get_width()/2, height),
                    xytext=(0, 3), textcoords="offset points",
                    ha='center', va='bottom', fontsize=10)

# ========== Subplot 2: Aggregation Efficiency (Fig 7b) ==========
clients = np.arange(10, 110, 10)
fedatt_time = 0.02 * clients**2  # FedAtt O(K²)
fedentgate_time = 0.38 * clients  # FedEntGate O(K)

# Plot lines

ax2.plot(clients, fedatt_time, 's-', label='FedAtt (O(K^2))', linewidth=2, markersize=6)
ax2.plot(clients, fedentgate_time, 'o-', label='FedEntGate (O(K))', linewidth=2, markersize=6)

# Set labels and title
ax2.set_xlabel('Number of Clients (K)', fontsize=12)
ax2.set_ylabel('Aggregation Time (ms)', fontsize=12)
ax2.set_title('(b) Aggregation Algorithm Complexity', fontsize=14,y=-0.3)
ax2.grid(linestyle='--', alpha=0.7)
ax2.legend(fontsize=12)

# Annotate key points
# ax2.annotate('K=100: 38ms', xy=(100, 38), xytext=(70, 100),
#              arrowprops=dict(arrowstyle='->', connectionstyle='arc3'),
#              fontsize=10)
# ax2.annotate('K=100: 210ms', xy=(100, 210), xytext=(70, 250),
#              arrowprops=dict(arrowstyle='->', connectionstyle='arc3'),
#              fontsize=10)

# Adjust layout and save
plt.tight_layout(pad=3.0)
plt.savefig('combined_figure.png', dpi=600)
plt.show()
